from numpy.core.fromnumeric import resize
from FourierTransform import *
from gauseKernel import *


class Conv2DWithUpsampling(Conv2D):
    def __init__(self, filters, kernel_size, sigma, P, padding="valid"):
        super().__init__(filters, kernel_size, sigma, boundary = "zero", padding=padding)
        self.P = P
    
    def multiply_element_wide(self, i, j, h, w, image, kernel):
        
        SUM = 0.0
        k = self.kernel_size // 2
        center = k
        tmpKernel = np.array(kernel)
        
        for m in range(-k, k+1):
            for n in range(-k, k+1):
                if i - m < 0 or j - n < 0 or i - m >= h or j - n >= w:
                    tmpKernel[center + m][center + n] = 0.0
                elif (i - m) % self.P != 0 or (j - n) % self.P != 0:
                    tmpKernel[center + m][center + n] = 0.0
        tmpSUM = 0.001
        for u in range(len(tmpKernel)):
            for t in range(len(tmpKernel)):
                tmpSUM += tmpKernel[u][t]
        tmpKernel = tmpKernel / tmpSUM
        
        for m in range(-k, k+1):
            for n in range(-k, k+1):
                if i - m >= 0 and j - n >= 0 and i - m < h and j - n < w:
                    tmp = tmpKernel[center + m][center + n] * image[h + i - m][w + j - n]
                    SUM += tmp
        return SUM

def downSampling(filepath, P, name):
    filepath = "./Resizing/" + filepath + ".jpg"
    imageRGB = Image.open(filepath)
    imageL = imageRGB.convert("L")
    imageL.save("rawImage.jpg")
    imageL_arr = np.array(imageL)
    print(imageL_arr.shape)
    conv = Conv2D(1, 3, 1, boundary = "zero", padding = "SAME")
    imageL_arr = conv(imageL_arr)[0]
    tmpImg = []
    for i in range(0, len(imageL_arr), P):
        tmpVector = []
        for j in range(0, len(imageL_arr[0]), P):
            tmpVector.append(imageL_arr[i][j])
        tmpImg.append(tmpVector)
    tmpImg = np.array(tmpImg)
    reconImg = Image.fromarray(tmpImg.astype("uint")).convert("L")
    nameALL = "./Resizing/" + name + ".jpg"
    reconImg.save(nameALL)
    # FT = FourierTransform()
    # B, F, Fabs, e = FT(imageL_arr)
    # Bcontainer = {}
    # for key, value in B.items():
    #     tmpB = []
    #     for i in range(0, len(value), P):
    #         tmpVector = []
    #         for j in range(0, len(value[0]), P):
    #             tmpVector.append(value[i][j])
    #         tmpB.append(tmpVector)
    #     Bcontainer[key] = np.array(tmpB)
    # return Bcontainer, F

def upSampling(filepath, P, name):
    filepath = "./Resizing/" + filepath + ".jpg"
    imageRGB = Image.open(filepath)
    imageL = imageRGB.convert("L")
    imageL_arr = np.array(imageL)
    print(imageL_arr.shape)
    conv = Conv2DWithUpsampling(1, 3, 1, P = 2, padding = "SAME")
    tmpImg = []
    for i in range(0, len(imageL_arr)):
        tmpVector = []
        for j in range(0, len(imageL_arr[0])):
            tmpVector.append(imageL_arr[i][j])
            for k in range(P-1):
                tmpVector.append(0)
        tmpImg.append(tmpVector)
        for l in range(P-1):
            tmpZeroVector = [0 for i in range(len(tmpVector))]
            tmpImg.append(tmpZeroVector)
    resized = Image.fromarray(np.array(tmpImg).astype("uint")).convert("L")
    #resized.show()
    resizedArr = np.array(conv(tmpImg)[0])
    resizeImage = Image.fromarray(resizedArr.astype("uint")).convert("L")
    nameALL = "./Resizing/" + name + ".jpg"
    resizeImage.save(nameALL)


def reconImage(B, F):
    h = len(B["0,0"])
    w = len(B["0,0"][0])
    print(h,w)
    tmpList = [[0.0 + 0.0j for i in range(w)] for j in range(h)]
    recon = np.array(tmpList)
    for i in range(h):
        for j in range(w):
            i_ = i - h // 2
            j_ = j - w // 2
            recon = recon + F[i][j] * F.shape[0] * F.shape[1] / (h * w) * B[str(i_) + "," + str(j_)]
    reconImage = Image.fromarray(recon.astype("uint")).convert("L")
    reconImage.save("./reconImage.jpg")


def getSub(fromReduce, fromExpand):
    reducedImage = Image.open("./Resizing/" + fromReduce + ".jpg")
    reducedInputs = np.array(reducedImage)
    print(np.max(reducedInputs))
    expandImage = Image.open("./Resizing/" + fromExpand + ".jpg")
    expandInputs = np.array(expandImage)
    print(np.min(expandInputs))
    if reducedInputs.shape != expandInputs.shape:
        raise Exception("两图尺寸不一致")
    sub1 = np.zeros(shape = reducedInputs.shape)
    for i in range(reducedInputs.shape[0]):
        for j in range(reducedInputs.shape[1]):
            sub1[i][j] = int(reducedInputs[i][j]) - int(expandInputs[i][j])
    print(np.max(sub1))
    subImage = Image.fromarray(sub1).convert("L")
    subImage.save("./Resizing/" + fromReduce + "-" + fromExpand + ".jpg")
    return sub1

filepath = "rawImage"
downSampled = "downSamplingImage"
downSampling(filepath, 2, downSampled)
upSampling(downSampled, 2, "upAfterDown")
getSub("rawImage", "upAfterDown")